# -*- coding: utf-8 -*-
from PyQt5.QtCore import *
from PyQt5.QtWidgets import *
from PyQt5.QtGui import *
import cv2
from src.gui.parser import FrameParser
import os
from glob import glob
import traceback

class Camera(object):

    def __init__(self):
        self.cap = cv2.VideoCapture(0)

    def read(self):
        return self.cap.read()

    def close(self):
        self.cap.release()


class MainWindow(QMainWindow):

    def __init__(self, parent=None):
        super(MainWindow, self).__init__(parent)
        # self.setupUi(self)
        self.setup_menu_bar()
        self.label = QLabel('', self)
        self.image_width = 640
        self.image_height = 360
        self.label.setFixedWidth(self.image_width)
        self.label.setFixedHeight(self.image_height)
        self.setCentralWidget(self.label)

        self._flag_camera_open = False
        self._flag_image_file = False

        self.frame_parser = FrameParser()

    def setup_menu_bar(self):
        exit_action = QAction('Exit', self)
        exit_action.triggered.connect(self.close)
        open_camera_action = QAction('Start!', self)
        open_camera_action.triggered.connect(self.start_camera)
        close_camera_action = QAction('Close', self)
        close_camera_action.triggered.connect(self.close_camera)
        open_folder_action = QAction('Open Folder..', self)
        open_folder_action.triggered.connect(self.open_folder)

        menu_bar = self.menuBar()
        menu_bar.setNativeMenuBar(False)

        file_menu = menu_bar.addMenu('&File')
        file_menu.addAction(open_folder_action)

        camera_menu = menu_bar.addMenu('&Camera')
        camera_menu.addAction(open_camera_action)
        camera_menu.addAction(close_camera_action)

        system_menu = menu_bar.addMenu('&System')
        system_menu.addAction(exit_action)

    def start_image_file(self):
        if self._flag_image_file is False:
            self._flag_image_file = True
            self.timer_image = QTimer(self)
            self.timer_image.timeout.connect(self.process_image_file)
            self.timer_image.start(0)

    def process_image_file(self):
        if self._flag_image_file:
            try:
                image_file = self.image_files[self.image_p]
                frame = cv2.imread(image_file)
                frame = cv2.resize(frame, (self.image_width, self.image_height))
                frame = self.frame_parser.parse(frame)
                frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
                showImage = QImage(frame.data, frame.shape[1], frame.shape[0], QImage.Format_RGB888)
                self.label.setPixmap(QPixmap.fromImage(showImage))
                self.image_p = (self.image_p + 1) % len(self.image_files)
                self.timer_image.start(16)
            except Exception:
                traceback.print_exc()
                self.image_p = (self.image_p + 1) % len(self.image_files)
                self.timer_image.start(16)


    def close_image_file(self):
        if self._flag_image_file:
            self._flag_image_file = False
            self.label.clear()

    def start_camera(self):
        self.close_image_file()
        if self._flag_camera_open is False:
            self.camera = Camera()
            self.timer_camera = QTimer(self)
            self.timer_camera.timeout.connect(self.process_camera)
            self._flag_camera_open = True
            self.timer_camera.start(0)

    def process_camera(self):
        if self._flag_camera_open:
            success, frame = self.camera.read()
            if success:
                frame = cv2.resize(frame, (self.image_width, self.image_height))
                frame = self.frame_parser.parse(frame)
                frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
                showImage = QImage(frame.data, frame.shape[1], frame.shape[0], QImage.Format_RGB888)
                self.label.setPixmap(QPixmap.fromImage(showImage))
                self.timer_camera.start(10)

    def close_camera(self):
        if self._flag_camera_open:
            self._flag_camera_open = False
            self.camera.close()
            self.label.clear()

    def open_folder(self):
        self.close_camera()
        self.close_image_file()
        path = os.path.abspath('.')
        path = os.path.join(path, './origin_dataset/dataset_picture/')
        folder_name = QFileDialog.getExistingDirectory(self, 'Open file', path)
        images = glob(os.path.join(folder_name, '*.jpg'))
        if len(images) == 0:
            return
        nums = [int(os.path.basename(f).split('.')[0]) for f in images]
        sorted_images = [x[0] for x in sorted(zip(images, nums))]
        self.image_files = sorted_images
        self.image_p = 0
        self.start_image_file()

if __name__ == '__main__':
    import sys
    app = QApplication(sys.argv)
    window = MainWindow()
    window.show()
    sys.exit(app.exec_())
